import re
import os
import argparse
from collections import defaultdict

# 预定义的正则表达式模式
PATTERNS = {
    'email': r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',
    'phone': r'\b(?:\+?86)?1[3-9]\d{9}\b',  # 中国手机号
    'url': r'\bhttps?://(?:[-\w.]|(?:%[\da-fA-F]{2}))+\.[a-zA-Z]{2,6}\b',
    'ip': r'\b(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\b',
    'idcard': r'\b[1-9]\d{5}(?:18|19|20)\d{2}(?:0[1-9]|1[0-2])(?:0[1-9]|[12]\d|3[01])\d{3}[\dXx]\b'  # 身份证号
}

def extract_info(text, patterns=None):
    """从文本中提取信息"""
    if not patterns:
        patterns = PATTERNS
    
    results = defaultdict(list)
    
    for name, pattern in patterns.items():
        matches = re.findall(pattern, text)
        # 去重并保持顺序
        unique_matches = []
        seen = set()
        for match in matches:
            if match not in seen:
                seen.add(match)
                unique_matches.append(match)
        results[name] = unique_matches
    
    return results

def process_file(file_path, output_file=None, extract_types=None):
    """处理单个文件"""
    try:
        # 读取文件内容
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            text = f.read()
        
        # 确定要提取的类型
        patterns = {}
        if extract_types:
            for type_name in extract_types:
                if type_name in PATTERNS:
                    patterns[type_name] = PATTERNS[type_name]
                else:
                    print(f"警告: 不支持的提取类型 '{type_name}'")
        else:
            patterns = PATTERNS
        
        # 提取信息
        results = extract_info(text, patterns)
        
        # 生成报告
        report = [f"文件: {file_path}"]
        total = 0
        
        for name, items in results.items():
            count = len(items)
            total += count
            report.append(f"\n{name} ({count}个):")
            for item in items:
                report.append(f"  - {item}")
        
        report_str = '\n'.join(report)
        print(report_str)
        
        # 保存结果
        if output_file:
            with open(output_file, 'a', encoding='utf-8') as f:
                f.write(report_str + '\n\n' + '-'*50 + '\n\n')
        
        return total
        
    except Exception as e:
        print(f"处理文件 {file_path} 失败: {str(e)}")
        return 0

def batch_process(files_dir, output_file=None, extract_types=None, recursive=False):
    """批量处理目录中的文件"""
    try:
        total_extracted = 0
        file_count = 0
        
        # 遍历目录
        for root, _, files in os.walk(files_dir):
            for file in files:
                # 只处理文本文件
                if file.lower().endswith(('.txt', '.html', '.htm', '.md', '.csv', '.log')):
                    file_path = os.path.join(root, file)
                    print(f"\n处理: {file_path}")
                    count = process_file(file_path, output_file, extract_types)
                    total_extracted += count
                    file_count += 1
            
            if not recursive:
                break
        
        print(f"\n批量处理完成！共处理 {file_count} 个文件，提取 {total_extracted} 条信息")
        
    except Exception as e:
        print(f"批量处理失败: {str(e)}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='文本内容提取脚本')
    parser.add_argument('path', help='文件路径或目录')
    parser.add_argument('--output', help='结果输出文件')
    parser.add_argument('--types', help='要提取的类型，用逗号分隔，如email,phone')
    parser.add_argument('--recursive', action='store_true', help='递归处理子目录')
    
    args = parser.parse_args()
    
    # 解析提取类型
    extract_types = args.types.split(',') if args.types else None
    
    # 处理单个文件或目录
    if os.path.isfile(args.path):
        process_file(args.path, args.output, extract_types)
    elif os.path.isdir(args.path):
        batch_process(args.path, args.output, extract_types, args.recursive)
    else:
        print(f"路径不存在: {args.path}")
